"""
detect_model.py
"""
import os
import threading

from ultralytics import YOLO


# 定义线程局部存储的类，用于保存模型实例
class ModelThreadLocal(threading.local):
    def __init__(self):
        self.model = None


# 创建线程局部存储对象
model_thread_local = ModelThreadLocal()


# 加载模型的函数
def load_model(detect_model):
    # 获取代码文件所在目录
    current_dir = os.path.dirname(os.path.abspath(__file__))

    # 项目根目录是代码文件目录的父目录
    project_dir = os.path.dirname(os.path.dirname(current_dir))

    # model文件路径
    model_path = os.path.join(project_dir, 'model', 'detect', detect_model)
    model = YOLO(model_path)
    return model


# 获取当前线程的模型实例
def get_model(detect_model='pms-slider-detect-20250518.pt'):
    # 检查当前线程是否有模型实例，如果没有则创建
    print(f"detect_model: {detect_model}")
    if not hasattr(model_thread_local, "model") or model_thread_local.model is None:
        model_thread_local.model = load_model(detect_model)
    return model_thread_local.model
